import traceback

import pandas as pd
import joblib
import torch

from utils.time_sequence import predict_time_range, add_time_features
from pathformer.PathFormer import PathFormer  # 确保路径正确
import utils.printer as printer

# 全局版本变量
version = "v2.0.0"

# 提取文件路径进行管理
data_path = 'dataset/ZGGG.09.02.2023.09.02.2025.1.0.0.cn.utf8.00000000.xls'
encoder_dd_path = f'model_files/label_encoder/label_encoder_dd_{version}.pkl'
encoder_c_path = f'model_files/label_encoder/label_encoder_c_{version}.pkl'
scaler_path = f'model_files/data_scaler/data_scaler_{version}.pkl'
model_dict_path = f'model_files/model_dict/pathformer_dict_{version}.pkl'
model_checkpoint_path = f'model_files/model_checkpoint/pathformer_checkpoint_{version}.ckpt'


# 加载原始数据后执行相同预处理
# 加载数据
data = pd.read_excel(data_path, skiprows=6)
data.drop(columns=['ff10', 'WW', "W'W'"], inplace=True)
data['当地时间 广州(机场)'] = pd.to_datetime(data['当地时间 广州(机场)'], format='%d.%m.%Y %H:%M')
data.set_index('当地时间 广州(机场)', inplace=True)
data.sort_index(ascending=True, inplace=True)  # 明确指定升序


# 加载原始数据后执行相同预处理
original_data = pd.read_excel('dataset/Guangzhou.15.01.2024.15.01.2025.xls', skiprows=6)
original_data.drop(columns=['ff10', 'WW', "W'W'"], inplace=True)
original_data['当地时间 广州(机场)'] = pd.to_datetime(original_data['当地时间 广州(机场)'], format='%d.%m.%Y %H:%M')
original_data.set_index('当地时间 广州(机场)', inplace=True)

data = add_time_features(data)

# 加载编码器和标准化器

encoder_dd = joblib.load(f'model_files/label_encoder/label_encoder_dd_{version}.pkl')
encoder_c = joblib.load(f'model_files/label_encoder/label_encoder_c_{version}.pkl')
scaler = joblib.load(f'model_files/data_scaler/data_scaler_{version}.pkl')

# 执行与训练时相同的文本处理
data['VV'] = data['VV'].astype(str)     # 确保数据类型为字符串
data['VV'] = data['VV'].str.extract(r'(\d+\.\d+)').astype(float)
data['DD'] = encoder_dd.transform(data[['DD']])  # 使用双括号选择DataFrame
data['c'] = encoder_c.transform(data[['c']])
data.ffill(inplace=True)  # 向前填充
le_dd = joblib.load('model_files/label_encoder_dd.pkl')
le_c = joblib.load('model_files/label_encoder_c.pkl')
scaler = joblib.load('model_files/data_scaler.pkl')

# 执行与训练时相同的文本处理
original_data['VV'] = original_data['VV'].str.extract(r'(\d+\.\d+)').astype(float)
original_data['DD'] = le_dd.transform(original_data['DD'])  # 使用加载的编码器
original_data['c'] = le_c.transform(original_data['c'])
original_data.ffill(inplace=True)  # 向前填充


# --------------------------
# 加载模型（关键修改点）
# --------------------------
# 方法1：直接加载完整检查点（需传入模型参数）
# model = PathFormer.load_from_checkpoint(
#     checkpoint_path='model_files/pathformerv1.0.0.ckpt',
#     input_dim=original_data.shape[1],  # 直接使用当前数据的特征数
#     patch_size=15,
#     embed_dim=256,
#     num_heads=8,
#     num_layers=6,
#     output_dim=original_data.shape[1]
# )
# # 设置模型为评估模式
# model.eval()

# 方法2：如果训练时保存的是 state_dict，需手动加载
model = PathFormer(
    input_dim=15,
    patch_size=48,
    embed_dim=128,
    num_heads=8,
    num_layers=4,
    output_dim=15,
    ff_dim=512
)
model.load_state_dict(torch.load(f'model_files/model_dict/pathformer_dict_{version}.pkl'))  # 从检查点提取 state_dict
model.eval()


# 调用预测函数
start_time = '2025-2-8 23:30'  # 必须存在于 original_data.index 中
num_steps = 48
# 设置模型为评估模式
model.eval()

# 方法2：如果训练时保存的是 state_dict，需手动加载
# model = PathFormer(
#     input_dim=original_data.shape[1],
#     patch_size=5,
#     embed_dim=64,
#     num_heads=8,
#     num_layers=4,
#     output_dim=original_data.shape[1]
# )
# model.load_state_dict(torch.load('model_files/pathformer.ckpt')['state_dict'])  # 从检查点提取 state_dict

# 调用预测函数
start_time = '2025-1-15 23:30'  # 必须存在于 original_data.index 中
num_steps = 50

try:
    pred_results = predict_time_range(
        model=model,
        data_scaler=scaler,
        data=data,
        start_time=start_time,
        num_steps=num_steps
    )

    # 反编码并输出
    pred_results['DD'] = encoder_dd.inverse_transform(pred_results[['DD']].astype(int)).flatten()
    # 将 Series 转换为 DataFrame
    pred_results['c'] = encoder_c.inverse_transform(pred_results[['c']].astype(int)).flatten()
    print(pred_results.to_string(max_rows=300, max_cols=None))
    pred_results['DD'] = le_dd.inverse_transform(pred_results['DD'].astype(int))
    pred_results['c'] = le_c.inverse_transform(pred_results['c'].astype(int))
    print(pred_results.to_string())
except ValueError as e:
    print(f"预测失败：{str(e)}")
    traceback.print_exc()
